import os
import random
import shutil
import cv2
from numpy.lib.type_check import imag

with open('data/All_data.txt','r') as f:
    lines = f.readlines()

train_percent = 0.9
num = len(lines)
print("total images numbers:", num)  

list = range(num)
tr = int(num * train_percent)
train = random.sample(list, tr)
print("train size:", tr)

ftrain = open('./data/train.txt', 'w')
fval = open('./data/val.txt', 'w')

for i in list:
    name = lines[i]
    line = name.strip().split(',')
    fn = line[0].split('\\')[-1].replace("json","jpg")
    xywh = [int(x) for x in line[1].split(' ')]
    x0,y0,w,h =  xywh[0]-20, xywh[1]-20,xywh[2]+40, xywh[3]+40       #w,h外扩40，是为了扩大切割区域，方便下步box数值计算
    xywh = [20,20,xywh[2],xywh[3]]
    pts = [int(x) for x in line[2:]]

    for j in range(len(pts)):
        if j%2==0:
            pts[j] = pts[j] -x0
        else:
            pts[j] = pts[j] -y0
    
    img = cv2.imread('zengqiang/'+fn)
    if x0<60 or y0<60 or y0+h>1540 or x0+w>1540:
        print('Error'+','+str(x0)+','+str(y0)+','+str(y0+h)+','+str(x0+w)+','+fn)
        continue
    # roi = img[y0-20:y0+h+20,x0-20:x0+w+20]
    roi = img[y0:y0 + h, x0:x0 + w]
    out = fn + " " + (',').join([str(a) for a in xywh]) +' '+(',').join([str(a) for a in pts]) +'\n'
    
    if i in train:
        cv2.imwrite("data/images_crop/train/"+fn,roi)
        ftrain.write(out)
    else:
        cv2.imwrite("data/images_crop/val/"+fn,roi)
        # cv2.imwrite("data/images_crop/all/"+fn,roi)
        fval.write(out)

ftrain.close()
fval.close()
print("write finished!")
